# Import the packages needed for the examples included in this notebook
import numpy as np
import matplotlib.pyplot as plt
import control as ct
print("python-control", ct.__version__)

# example 1:Open Loop Analysis of a Coupled Mass Spring System
k=2
m=1
c=0.1
# dx/dt = Ax + Bu
# 
A = np.array([[0,0,1,0],
        [0,0,0,1],
        [-2*k/m, k/m, -c/m,.0],
        [k/m, -2*k/m, 0, -c/m]])
B = np.array([[0],
        [0],
        [0],
        [k/m]])

C = np.array([[1, 0, 0, 0],
        [0, 1, 0, 0]])

D = 0

sys = ct.ss(A, B, C, D, outputs=['q1', 'q2'], name='Coupled Spring Mass')
print(sys)

# initial conditions without input
response = ct.initial_response(sys, X0=[1, 0, 0, 0])
cplt = response.plot()
# Plot the outputs of the system on the same graph, in different colors
t = response.time
x = response.states
plt.plot(t, x[0], 'b', t, x[1], 'r')
plt.legend(['$x_1$', '$x_2$'])
plt.xlim(0, 50)
plt.ylabel('States')
plt.xlabel('Time [s]')
plt.title("Initial response from $x_1 = 1$, $x_2 = 0$")
plt.show()

# various initial conditions 无输入响应
for X0 in [[1, 0, 0, 0], [0, 2, 0, 0], [1, 2, 0, 0], [0, 0, 1, 0], [0, 0, 2, 0]]:
  response = ct.initial_response(sys, T=20, X0=X0)
  response.plot(label=f"{X0=}")

plt.show()

# step response 阶跃响应
stepresp = ct.step_response(sys)
cplt = stepresp.plot(plot_inputs=True)
plt.show()

# Forced Response 时变输入响应u(t)
T = np.linspace(0, 50, 100)
U1 = np.cos(T)
U2 = np.sin(3*T)
U3 = U1 + U2

resp1 = ct.forced_response(sysdata=sys, timepts=T, inputs=U1)
resp2 = ct.forced_response(sysdata=sys, timepts=T, inputs=U2)
resp3 = ct.forced_response(sysdata=sys, timepts=T, inputs=U3)

resp1.sysname = 'U1'
resp2.sysname = 'U2'
resp3.sysname = 'U3'

resp1.plot(color='r', label=f"{resp1.sysname}")
resp2.plot(color='b', label=f"{resp2.sysname}")
resp3.plot(color='g', label=f"{resp3.sysname}")
plt.show()

# 通过线性叠加输入和输出，判断系统是否是线性的
cplt = resp3.plot(label="G(U1+U2)")
cplt.axes[0,0].plot(resp1.time, resp1.outputs[0,:] + resp2.outputs[0,:], 'k--', label="G(U1)+G(U2)")
cplt.axes[1,0].plot(resp1.time, resp1.outputs[1,:] + resp2.outputs[1,:], 'k--')
cplt.axes[2,0].plot(resp1.time, resp1.inputs[0,:] + resp2.inputs[0,:], 'k--', label="U1+U2")
plt.show()

# forced inputs from non-zero initial conditions is not linear
X0=[1,0,0,0]
resp1 = ct.forced_response(sys, T, U1, X0)
resp2 = ct.forced_response(sys, T, U2, X0)
resp3 = ct.forced_response(sys, T, U1+U2, X0)

cplt = resp3.plot(label='G(U1 + U2)')
cplt.axes[0, 0].plot(resp1.time, resp1.outputs[0,:] )